import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from imblearn.over_sampling import SMOTE
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from xgboost import XGBClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, roc_curve, auc
import os
from sklearn.tree import DecisionTreeClassifier
from sklearn.svm import SVC

# --- 全局绘图设置 ---
sns.set(style="whitegrid", palette="pastel")
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def plot_confusion_matrix(y_true, y_pred, classes, save_path, title):
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(5, 4))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=classes, yticklabels=classes)
    plt.xlabel('预测值')
    plt.ylabel('真实值')
    plt.title(title)
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()


def plot_roc_curve(y_true, y_score, save_path, title):
    plt.figure(figsize=(6, 5))
    fpr, tpr, _ = roc_curve(y_true, y_score)
    roc_auc = auc(fpr, tpr)
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=1, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title(title)
    plt.legend(loc="lower right")
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()


def safe_drop(df, columns):
    return df.drop(columns=[col for col in columns if col in df.columns], errors='ignore')


def load_data():
    # 读取数据
    stroke = pd.read_csv('../preprogressData/output/stroke_processed.csv')
    heart = pd.read_csv('../preprogressData/output/heart_processed.csv')
    cirrhosis = pd.read_csv('../preprogressData/output/cirrhosis_processed.csv')
    return stroke, heart, cirrhosis


def preprocess_stroke(df):
    # 只保留有用特征，编码类别变量
    X = safe_drop(df, ['id', 'stroke'])
    y = df['stroke']
    X = pd.get_dummies(X)
    return X, y


def preprocess_heart(df):
    X = safe_drop(df, ['HeartDisease'])
    y = df['HeartDisease']
    X = pd.get_dummies(X)
    return X, y


def preprocess_cirrhosis(df, binary=True):
    # 剔除不需要的列
    X = safe_drop(df, ['ID', 'N_Days', 'Status', 'Stage'])
    y = df['Stage']
    if binary:
        y = (y >= 3).astype(int)
    X = pd.get_dummies(X)
    return X, y


def train_and_evaluate(X, y, disease_name, save_dir):
    os.makedirs(save_dir, exist_ok=True)

    # 划分训练集和测试集 (80%训练, 20%测试)
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=42, stratify=y
    )

    # 特征缩放
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # 对中风数据应用SMOTE重采样（仅用于二分类问题）
    if disease_name == 'stroke' or disease_name == 'heart':
        smote = SMOTE(random_state=42)
        X_train_res, y_train_res = smote.fit_resample(X_train_scaled, y_train)

        print(f"\n{disease_name} - SMOTE处理前训练集类别分布:")
        print(pd.Series(y_train).value_counts())
        print(f"\n{disease_name} - SMOTE处理后训练集类别分布:")
        print(pd.Series(y_train_res).value_counts())
    else:
        # 对于肝硬化数据，不使用重采样
        X_train_res, y_train_res = X_train_scaled, y_train

    # 定义所有模型
    models = {
        '逻辑回归': LogisticRegression(max_iter=1000, random_state=42),
        '随机森林': RandomForestClassifier(random_state=42),
        'XGBoost': XGBClassifier(random_state=42, eval_metric='logloss'),
        '决策树': DecisionTreeClassifier(random_state=42),
        'SVM': SVC(probability=True, random_state=42)
    }

    # 存储结果
    all_results = {}

    for name, model in models.items():
        print(f"\n{'-' * 20} 训练模型: {disease_name} - {name} {'-' * 20}")

        try:
            # 训练模型
            model.fit(X_train_res, y_train_res)

            # 预测
            y_pred = model.predict(X_test_scaled)

            # 分类报告
            report = classification_report(y_test, y_pred, output_dict=True)
            print(f"\n{disease_name} - {name} 分类报告:")
            print(classification_report(y_test, y_pred, digits=4))

            # 保存报告
            pd.DataFrame(report).to_csv(os.path.join(save_dir, f'{disease_name}_{name}_report.csv'))

            # 混淆矩阵
            class_names = ['健康', '患病'] if disease_name in ['stroke', 'heart'] else sorted(y.unique().astype(str))
            plot_confusion_matrix(
                y_test, y_pred, class_names,
                os.path.join(save_dir, f'{disease_name}_{name}_cm.png'),
                f'{disease_name}-{name} 混淆矩阵'
            )

            # ROC曲线（仅二分类）
            if disease_name in ['stroke', 'heart']:
                y_score = model.predict_proba(X_test_scaled)[:, 1]
                auc_score = roc_auc_score(y_test, y_score)
                print(f"AUC分数: {auc_score:.4f}")

                plot_roc_curve(
                    y_test, y_score,
                    os.path.join(save_dir, f'{disease_name}_{name}_roc.png'),
                    f'{disease_name}-{name} ROC曲线 (AUC={auc_score:.2f})'
                )

                # 存储结果
                all_results[name] = {
                    'classification_report': report,
                    'auc': auc_score
                }

        except Exception as e:
            print(f"训练模型 {name} 时出错: {str(e)}")

    return all_results


def main():
    # 创建输出目录
    os.makedirs('prediction_results/stroke', exist_ok=True)
    os.makedirs('prediction_results/heart', exist_ok=True)
    os.makedirs('prediction_results/cirrhosis', exist_ok=True)

    # 加载数据
    stroke, heart, cirrhosis = load_data()

    # 处理中风数据
    print("\n" + "=" * 50)
    print("开始处理中风数据")
    print("=" * 50)
    X_s, y_s = preprocess_stroke(stroke)
    stroke_results = train_and_evaluate(X_s, y_s, 'stroke', 'prediction_results/stroke')

    # 处理心脏病数据
    print("\n" + "=" * 50)
    print("开始处理心脏病数据")
    print("=" * 50)
    X_h, y_h = preprocess_heart(heart)
    heart_results = train_and_evaluate(X_h, y_h, 'heart', 'prediction_results/heart')

    # 处理肝硬化数据
    print("\n" + "=" * 50)
    print("开始处理肝硬化数据")
    print("=" * 50)
    X_c, y_c = preprocess_cirrhosis(cirrhosis, binary=True)
    cirrhosis_results = train_and_evaluate(X_c, y_c, 'cirrhosis', 'prediction_results/cirrhosis')

    # 输出最佳模型
    print("\n" + "=" * 50)
    print("最佳模型总结")
    print("=" * 50)

    for disease, results in zip(['中风', '心脏病', '肝硬化'], [stroke_results, heart_results, cirrhosis_results]):
        if results:
            best_model = max(results, key=lambda k: results[k]['auc'] if 'auc' in results[k] else 0)
            best_auc = results[best_model].get('auc', 'N/A')
            print(f"{disease}最佳模型: {best_model} (AUC = {best_auc})")


if __name__ == '__main__':
    main()